import json
import os
from py2neo import Graph, Node, Relationship
from tqdm import tqdm

CONFIG = {
    "neo4j_uri": "bolt://localhost:7687",
    "neo4j_auth": ("neo4j", "12345678"),
    "data_path": os.path.join('.', 'data.json'),
    "batch_size": 10
}

def get_graph_connection():
    """获取Neo4j数据库连接"""
    return Graph(CONFIG["neo4j_uri"], auth=CONFIG["neo4j_auth"])

def validate_data_path():
    """验证数据文件路径"""
    abs_path = os.path.abspath(CONFIG["data_path"])
    if not os.path.exists(abs_path):
        raise FileNotFoundError(f"数据文件不存在: {abs_path}")

def load_data():
    """加载JSON数据"""
    with open(CONFIG["data_path"], 'r', encoding='utf-8') as file:
        return [json.loads(line) for line in file]

def create_indexes(graph):
    """创建必要的索引"""
    indexes = [
        "CREATE INDEX IF NOT EXISTS FOR (n:方剂) ON (n.name)",
        "CREATE INDEX IF NOT EXISTS FOR (n:药材) ON (n.name)",
        "CREATE INDEX IF NOT EXISTS FOR (n:疾病) ON (n.name)",
        "CREATE INDEX IF NOT EXISTS FOR (n:症状) ON (n.name)"
    ]
    for query in indexes:
        graph.run(query)

def process_item(tx, item):
    """处理单个条目"""
    if "name" not in item or not item["name"]:
        return

    # 创建/合并方剂节点
    方剂 = Node("方剂", name=item["name"],
               consist=item.get("consist", ""),
               use=item.get("use", ""))
    tx.merge(方剂, "方剂", "name")

    # 处理药材关系
    for herb in item.get("formual", []):
        if herb:
            药材 = Node("药材", name=herb)
            tx.merge(药材, "药材", "name")
            tx.create(Relationship(方剂, "包含药材", 药材))

    # 处理疾病关系
    for disease in item.get("sick", []):
        if disease:
            疾病 = Node("疾病", name=disease)
            tx.merge(疾病, "疾病", "name")
            tx.create(Relationship(方剂, "主治疾病", 疾病))

    # 处理症状关系
    for symptom in item.get("effect", []):
        if symptom:
            症状 = Node("症状", name=symptom)
            tx.merge(症状, "症状", "name")
            tx.create(Relationship(方剂, "缓解症状", 症状))

def process_batch(graph, batch):
    """处理数据批次"""
    tx = graph.begin()
    try:
        for item in batch:
            process_item(tx, item)
        tx.commit()
    except Exception as e:
        tx.rollback()
        raise RuntimeError(f"处理批次时出错: {str(e)}")

def import_data_to_neo4j():
    """主数据导入函数"""
    try:
        print(f"当前工作目录: {os.getcwd()}")
        validate_data_path()
        data = load_data()
        graph = get_graph_connection()

        create_indexes(graph)

        valid_data = [item for item in data if "name" in item and item["name"]]
        print(f"总数据量: {len(data)}, 有效数据量: {len(valid_data)}")

        batches = [
            valid_data[i:i + CONFIG["batch_size"]]
            for i in range(0, len(valid_data), CONFIG["batch_size"])
        ]

        for batch in tqdm(batches, desc="导入进度"):
            process_batch(graph, batch)

        print(f"成功导入 {len(valid_data)} 条方剂数据")
    except Exception as e:
        print(f"导入失败: {str(e)}")
        raise

if __name__ == "__main__":
    import_data_to_neo4j()